#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# !/bin/env python
from __future__ import unicode_literals

from concurrent import futures

import grpc

from proto import recall_pb2 
from proto import recall_pb2_grpc 
from proto import user_info_pb2 as user_info_pb2
import redis
from milvus import Milvus, DataType
from paddle_serving_app.local_predict import LocalPredictor
import numpy as np


from milvus_tool.milvus_recall import RecallByMilvus

def hash2(a):
    return hash(a) % 20000000
#实现服务的接口
class RecallServerServicer(object):
    def __init__(self):
        #加载用户模型，用户用户特征提取
        self.uv_client = LocalPredictor()
        self.uv_client.load_model_config("user_vector_model/serving_server_dir",use_gpu=True) 
        # milvus_host = '127.0.0.1'
        # milvus_port = '19530'
        # self.milvus_client = Milvus(milvus_host, milvus_port)
        #声明milvus对象，以及要查找的集合名称
        self.milvus_client = RecallByMilvus()
        self.collection_name = 'demo_e_commerce'
    #用户特征向量获取
    def get_user_vector(self, user_info):
        '''
        um_res.user_info.user_id = user_info["user_id"]
        um_res.user_info.age = user_info["age"]
        um_res.user_info.sex = user_info["sex"]
        um_res.user_info.city_level = user_info["city_level"]
        um_res.user_info.province = user_info["province"]
        um_res.user_info.city = user_info["city"]
        um_res.user_info.country = user_info["country"]
        '''
        dic = {"userid": [], "age": [], "sex": [], "city_level": [], "province": [], "city": [], "country": []}
        lod = [0]
        dic["userid"].append(hash2(user_info.user_id))
        dic["age"].append(hash2(user_info.age))
        dic["sex"].append(hash2(user_info.sex))
        dic["city_level"].append(hash2(user_info.city_level))
        dic["province"].append(hash2(user_info.province))
        dic["city"].append(hash2(user_info.city))
        dic["country"].append(hash2(user_info.country))
        lod.append(1)

        dic["userid.lod"] = lod
        dic["age.lod"] = lod
        dic["sex.lod"] = lod
        dic["city_level.lod"] = lod
        dic["province.lod"] = lod
        dic["city.lod"] = lod
        dic["country.lod"] = lod
        for key in dic:
            dic[key] = np.array(dic[key]).astype(np.int64).reshape(len(dic[key]),1)
        #paddleServing模型推理服务，如果您更换了模型请在此处更改参数
        fetch_map = self.uv_client.predict(feed=dic, fetch=["save_infer_model/scale_0.tmp_7"], batch=True)
        return fetch_map["save_infer_model/scale_0.tmp_7"].tolist()[0]
    #召回服务，通过milvus，召回用户的候选集商品id列表
    def recall(self, request, context):
        '''
    message RecallRequest{
        string log_id = 1;
        user_info.UserInfo user_info = 2;
        string recall_type= 3;
        uint32 request_num= 4;
    }

    message RecallResponse{
        message Error {
            uint32 code = 1;
            string text = 2;
        }
        message ScorePair {
            string nid = 1;
            float score = 2;
        };
        Error error = 1;
        repeated ScorePair score_pairs = 2;
    }
        '''
        recall_res = recall_pb2.RecallResponse()
        user_vector = self.get_user_vector(request.user_info)
        #在milvus中召回与该用户向量相似的top_k个电影id，返回结果为id和distance（向量间的相似度）
        status, results = self.milvus_client.search(collection_name=self.collection_name, vectors=[user_vector])
        for entities in results:
            if len(entities) == 0:
                recall_res.error.code = 500
                recall_res.error.text = "Recall server get milvus fail. ({})".format(str(request))
                return recall_res
            for topk_prod in entities:
                current_entity = topk_prod.id
                score_pair = recall_res.score_pairs.add()
                score_pair.nid = str(topk_prod.id)
                score_pair.score = float(topk_prod.distance)
        recall_res.error.code = 200
        return recall_res
#定义服务
class RecallServer(object):
    """
    recall server
    """
    def start_server(self):
        max_workers = 40  # 定义多线程的服务器对象
        concurrency = 40  # 定义最大连接数量
        port = 8950  # 定义服务端口
        #开启服务，对外提供rpc调用
        server = grpc.server(
            futures.ThreadPoolExecutor(max_workers=max_workers),
            options=[('grpc.max_send_message_length', 1024 * 1024),
                     ('grpc.max_receive_message_length', 1024 * 1024)],
            maximum_concurrent_rpcs=concurrency)
        #注册实现服务的方法到服务器对象中
        servicer = RecallServerServicer()
        recall_pb2_grpc.add_RecallServiceServicer_to_server(servicer, server)
        #为服务绑定主机与端口
        server.add_insecure_port('[::]:{}'.format(port))
        #开启服务
        server.start()
        print('recall服务启动！')
        server.wait_for_termination()

if __name__ == "__main__":
    recall = RecallServer()
    recall.start_server()
